/*
 * Wazuh Vulnerability scanner - Scan Orchestrator
 * Copyright (C) 2015, Wazuh Inc.
 * Nov 23, 2023.
 *
 * This program is free software; you can redistribute it
 * and/or modify it under the terms of the GNU General Public
 * License (version 2) as published by the FSF - Free Software
 * Foundation.
 */

#ifndef _OS_DATA_CACHE_HPP
#define _OS_DATA_CACHE_HPP

#include "../policyManager/policyManager.hpp"
#include "cacheLRU.hpp"
#include "loggerHelper.h"
#include "singleton.hpp"
#include "socketDBWrapper.hpp"
#include "vulnerabilityScannerDefs.hpp"
#include "wazuhDBQueryBuilder.hpp"
#include <shared_mutex>
#include <string>

auto constexpr WDB_SOCKET {"queue/db/wdb"};

/**
 * @brief Os structure.
 */
struct Os final
{
    std::string hostName;       ///< Hostname of the agent.
    std::string architecture;   ///< Architecture of the agent.
    std::string name;           ///< Name of operating system.
    std::string codeName;       ///< Code name of the operating system.
    std::string majorVersion;   ///< Major version of the operating system.
    std::string minorVersion;   ///< Minor version of the operating system.
    std::string patch;          ///< Patch of the operating system.
    std::string build;          ///< Build of the operating system.
    std::string platform;       ///< Platform of the operating system.
    std::string version;        ///< Version of the operating system.
    std::string release;        ///< Release of the operating system.
    std::string displayVersion; ///< Display version of the operating system.
    std::string sysName;        ///< System name of the operating system.
    std::string kernelVersion;  ///< Version of the kernel operating system.
    std::string kernelRelease;  ///< Release of the kernel operating system.
    std::string cpeName;        ///< CPE name of the operating system.
};

/**
 * @brief OsDataException class.
 *
 */
class OsDataException : public std::exception
{
public:
    /**
     * @brief Overload what() method.
     *
     * @return const char*
     */
    // LCOV_EXCL_START
    const char* what() const noexcept override
    {
        return m_msg.what();
    }
    // LCOV_EXCL_STOP

    /**
     * @brief Construct a new Os Data Exception object
     *
     * @param message
     */
    OsDataException(const std::string& message)
        : m_msg {message}
    {
    }

private:
    std::runtime_error m_msg;
};

/**
 * @brief OsDataCache class.
 */
class OsDataCache final : public Singleton<OsDataCache>
{
private:
    LRUCache<std::string, Os> m_osData {PolicyManager::instance().getOsdataLRUSize()};
    std::shared_mutex m_mutex;
    std::optional<SocketDBWrapper> m_wdbSocketWrapper {std::nullopt};

    Os getOsDataFromWdb(const std::string& agentId)
    {
        nlohmann::json response;
        try
        {
            m_wdbSocketWrapper->query(WazuhDBQueryBuilder::builder().agentGetOsInfoCommand(agentId).build(), response);
        }
        catch (const std::exception& e)
        {
            throw OsDataException(e.what());
        }

        if (response.empty())
        {
            throw OsDataException("Empty response from wazuh-db (no OS data)");
        }

        Os osData;
        auto& data = response.at(0);

        osData.hostName = data.value("hostname", "");
        osData.architecture = data.value("architecture", "");
        osData.name = data.value("os_name", "");
        osData.codeName = data.value("os_codename", "");
        osData.majorVersion = data.value("os_major", "");
        osData.minorVersion = data.value("os_minor", "");
        osData.patch = data.value("os_patch", "");
        osData.build = data.value("os_build", "");
        osData.platform = data.value("os_platform", "");
        osData.version = data.value("os_version", "");
        osData.release = data.value("os_release", "");
        osData.displayVersion = data.value("os_display_version", "");
        osData.sysName = data.value("sysname", "");
        osData.kernelVersion = data.value("version", "");
        osData.kernelRelease = data.value("release", "");
        osData.cpeName = data.value("cpe_name", "");

        return osData;
    }

public:
    /**
     * @brief This method returns the os data.
     * @param agentId agent id.
     *
     * @return Os
     */
    Os getOsData(const std::string& agentId)
    {
        std::shared_lock<std::shared_mutex> lock(m_mutex);
        auto value = m_osData.getValue(agentId);

        if (value)
        {
            return *value;
        }

        if (!m_wdbSocketWrapper)
        {
            try
            {
                m_wdbSocketWrapper.emplace(WDB_SOCKET);
            }
            catch (...)
            {
                throw OsDataException("Error creating socketDBWrapper");
            }
        }

        // This may throw an exception that will be captured by the caller method.
        auto osData = getOsDataFromWdb(agentId);
        m_osData.insertKey(agentId, osData);
        return osData;
    }

    /**
     * @brief This method sets the os data.
     *
     * @param agentId agent id.
     * @param osData os data.
     */
    void setOsData(const std::string& agentId, const Os& osData)
    {
        std::lock_guard<std::shared_mutex> lock(m_mutex);
        m_osData.insertKey(agentId, osData);
    }
};
#endif // _OS_DATA_CACHE_HPP
